#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat May  8 03:08:38 2021

@author: sarupa
"""

#Imports 
import pandas as pd
import copy
import numpy as np
import tensorflow as tf
from tensorflow import keras
from matplotlib import pyplot as plt
from tensorflow.keras.models import Sequential
from keras.optimizers import Adam
from tensorflow.keras.layers import LSTM, Dense, Dropout, SimpleRNN,Flatten,TimeDistributed
from sklearn.preprocessing import MinMaxScaler
from LSTM_model import *
from casadi import *
from numpy import linalg as la
import control
import timeit
import time
from numpy import random
from random import seed

start = timeit.default_timer()

# Define RMSE
def rmse(predictions, targets):

    differences = (predictions - targets)/targets                      

    differences_squared = differences ** 2                    

    mean_of_differences_squared = differences_squared.mean()  

    rmse_val = np.sqrt(mean_of_differences_squared)*100         

    return rmse_val 

# Load the data
df=pd.read_csv("Final_data_set.csv") # This data has already been resampled

# Prepare the data
dataset_in=df.iloc[:,].values
dataset=np.zeros((len(dataset_in),10))
dataset[:,0]=dataset_in[:,7]
dataset[:,1]=dataset_in[:,4]
dataset[:,2]=dataset_in[:,2]
dataset[:,3]=dataset_in[:,5]
dataset[:,4:10]=dataset_in[:,9:15]


# The min-max scaling approach
sc= MinMaxScaler(feature_range=(0,1))
dataset_scaled = sc.fit_transform(dataset)

# Scale the test data using the scaling parameters of the original data
# dataset_scaled = (dataset - data_min)/(data_max - data_min)

# Extract the scaling parameters
data_min = sc.data_min_
data_max = sc.data_max_

# Split the original data into training, validation, and testing data sets
data_train = dataset_scaled[:200000,:]
data_validation  = dataset_scaled[400000:475000,:]
data_test = dataset_scaled[475000:480000,:]

# Prepare the data for the RNN
sequenceLength = 2
numberOfFeatures = 10
numberOfOutputs = 1
sequenceLength_out=2
number_of_states= 4

# Prepare the training data
X_train, y_train = [], []

for i in range(sequenceLength, len(data_train)-sequenceLength):
    X_train.append(data_train[i-sequenceLength:i,0:numberOfFeatures])
    y_train.append(data_train[i:i+numberOfOutputs,0:number_of_states])

X_train, y_train = np.array(X_train), np.array(y_train)

X_train = X_train.reshape(X_train.shape[0], X_train.shape[1], numberOfFeatures)
y_train = y_train.reshape(y_train.shape[0], y_train.shape[1], numberOfOutputs*number_of_states)

# Prepare the validation data
X_val, y_val = [], []

for i in range(sequenceLength, len(data_validation)-sequenceLength):
    
    X_val.append(data_validation[i-sequenceLength:i,0:numberOfFeatures])
    y_val.append(data_validation[i:i+numberOfOutputs,0:number_of_states])

X_val, y_val = np.array(X_val), np.array(y_val) 

X_val = X_val.reshape(X_val.shape[0], X_val.shape[1], numberOfFeatures)
y_val = y_val.reshape(y_val.shape[0], y_val.shape[1], numberOfOutputs*number_of_states)

# Training the model
model=Sequential()
model.add(LSTM(units=50, input_shape=(sequenceLength,numberOfFeatures), return_sequences=True))
model.add(LSTM(units=50, return_sequences=True))
model.add(TimeDistributed(Dense(numberOfOutputs*number_of_states)))
model.compile(loss='mse', optimizer='adam', metrics='mean_squared_error')

history = model.fit(X_train, y_train, epochs=30, batch_size=100, validation_data=(X_val, y_val))#, callbacks=[callbacks])


# Prepare the test data for validation of the RNN model

X_test, y_test = [], []

for i in range(sequenceLength, len(data_test)-sequenceLength):
    
    X_test.append(data_test[i-sequenceLength:i,0:numberOfFeatures])
    y_test.append(data_test[i:i+numberOfOutputs,0:number_of_states])

X_test, y_test = np.array(X_test), np.array(y_test)

X_test = X_test.reshape(X_test.shape[0], X_test.shape[1], numberOfFeatures)
y_test = y_test.reshape(y_test.shape[0], y_test.shape[1], numberOfOutputs*number_of_states)

# Prediction of test data
predicted_trajectory=model.predict(X_test)
predicted_trajectory=predicted_trajectory[:,-1,:]
y_test = y_test.reshape(y_test.shape[0],-1)

# Rescale to original data
y_test_unscaled = y_test*(data_max[0:number_of_states]-data_min[0:number_of_states]) + data_min[0:number_of_states]
predicted_trajectory_unscaled = predicted_trajectory*(data_max[0:number_of_states]-data_min[0:number_of_states]) + data_min[0:number_of_states]

# Data plot 
plt.figure(1)
str_output = ['CB3' ,'CB2','T1 [K]','T2 [K]']

for i in range(number_of_states): 
    plt.subplot(number_of_states,1,i+1)
    plt.plot(y_test_unscaled[0:400,i],'m')
    plt.plot(predicted_trajectory_unscaled[0:400,i],'b')
    plt.ylabel("Tempstep " + str(i+1))
    plt.ylabel(str_output[i])
    plt.legend(['Data','RNN'])
    plt.grid()
    plt.xlabel("Timestep")
    if i == 0:
        plt.title('PredictedTrajectory_Rnn Vs Actual trajectory')
        #plt.savefig("multiple_test4_result1")     

plt.figure(2)
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('Model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['train', 'test'], loc='upper left')
plt.show()

print("MSE of model:%.2f"% mean_squared_error(predicted_trajectory_unscaled, y_test_unscaled))

# Extract the weights and obtain the mathematical expression of the LSTM [ START ]

# Layer 1
# The number of units
units_1 = int(model.layers[0].trainable_weights[0].shape[1]/4)

W_1=model.layers[0].get_weights()[0] # The Kernel, the weighting matrix for the input
U_1=model.layers[0].get_weights()[1] # The recurrent kernel, the weighting matrix for the hidden states
b_1=model.layers[0].get_weights()[2] # The bias matrix

# Obtain the matrices for the gates and the cell state

# The kernel matrices for the gates and the cell state
W_i_1 = W_1[:,:units_1] # The input gate
W_f_1 = W_1[:,units_1: units_1*2] # The forget gate
W_c_1 = W_1[:, units_1*2:units_1*3] # The cell state
W_o_1 = W_1[:, units_1*3:] # The output gate

# The recurrent kernel matrices for the gates and the cell state

U_i_1 = U_1[:,:units_1] # The input gate
U_f_1 = U_1[:,units_1:units_1*2] # The forget gate
U_c_1 = U_1[:, units_1*2:units_1*3] # The cell state
U_o_1 = U_1[:, units_1*3:] # The output gate


# The bias matrices for the gates and the cell state

b_i_1 = b_1[:units_1] # The input gate
b_f_1 = b_1[units_1:units_1*2] # The forget gate
b_c_1 = b_1[units_1*2:units_1*3] #The cell state
b_o_1 = b_1[units_1*3:] # The output gate

# Layer 2
# The number of units
units_2 = int(model.layers[1].trainable_weights[0].shape[1]/4)

W_2=model.layers[1].get_weights()[0] # The Kernel, the weighting matrix for the input
U_2=model.layers[1].get_weights()[1] # The recurrent kernel, the weighting matrix for the hidden states
b_2=model.layers[1].get_weights()[2] # The bias matrix

# Obtain the matrices for the gates and the cell state

#The kernel matrices for the gates and the cell state
W_i_2 = W_2[:,:units_2] # The input gate
W_f_2 = W_2[:,units_2: units_2*2] # The forget gate
W_c_2 = W_2[:, units_2*2:units_2*3] # The cell state
W_o_2 = W_2[:, units_2*3:] # The output gate

# The recurrent kernel matrices for the gates and the cell state

U_i_2 = U_2[:,:units_2] # The input gate
U_f_2 = U_2[:,units_2:units_2*2] # The forget gate
U_c_2 = U_2[:, units_2*2:units_2*3] # The cell state
U_o_2 = U_2[:, units_2*3:] # The output gate


# The bias matrices for the gates and the cell state

b_i_2 = b_2[:units_2] # The input gate
b_f_2 = b_2[units_2:units_2*2] # The forget gate
b_c_2 = b_2[units_2*2:units_2*3] #The cell state
b_o_2 = b_2[units_2*3:] # The output gate



# The Dense layer
V_m=model.layers[2].get_weights()[0]
b_m=model.layers[2].get_weights()[1]


# Extract the weights and obtain the mathematical expression of the LSTM [ END ]



# The mathematical expression for the LSTM and then compare the results with the model.predict 
# LSTM model is written explicitly in LSTM_model file
y_rnn_model= np.zeros((y_test.shape[0], numberOfOutputs*number_of_states))

for k in range(y_test.shape[0]):
    print(k)
    x = X_test[k,:,0:number_of_states]
    u = X_test[k,:,number_of_states:numberOfFeatures]
    y_rnn_model[k,:] = ObtainlstmModel_1(W_1, U_1, b_1, W_2, U_2, b_2, V_m, b_m, units_1, units_2, sequenceLength, x,u)
y_rnn_model_unscaled = y_rnn_model*(data_max[0:number_of_states]-data_min[0:number_of_states]) + data_min[0:number_of_states]


# Multiple time steps ahead prediction using the LSTM [ START ]

x_test_MSAP = copy.deepcopy(X_test)

x_test_2D = x_test_MSAP.reshape(x_test_MSAP.shape[0], -1)
    
x_test_2D = x_test_2D[:,0:numberOfFeatures]

initial_input=x_test_2D[0:sequenceLength]

x_in=initial_input[:,0:number_of_states]
u_in=initial_input[:,number_of_states:numberOfFeatures]

y_test_MSAP = []
y_test_MSAP.append(ObtainlstmModel_1(W_1, U_1, b_1, W_2, U_2, b_2, V_m, b_m, units_1, units_2, sequenceLength, x_in, u_in))

for i in range(1, len(x_test_2D)-sequenceLength+1):
    #print(i)
    current_input = x_test_2D[i:i+sequenceLength]
    current_input[-1,0:number_of_states] = y_test_MSAP[i-1]
    x_test_2D[i:i+sequenceLength] = current_input
    x_curr=current_input[:,0:number_of_states]
    u_curr=current_input[:,number_of_states:numberOfFeatures]
    y_test_MSAP.append(ObtainlstmModel_1(W_1, U_1, b_1, W_2, U_2, b_2, V_m, b_m, units_1, units_2, sequenceLength, x_curr,u_curr))

y_test_MSAP = np.array(y_test_MSAP)
y_test_MSAP_unscaled = y_test_MSAP*(data_max[0:number_of_states]-data_min[0:number_of_states]) + data_min[0:number_of_states]

plt.figure(3)
str_output = ['CB2' ,'CB3','T1 [K]', 'T2 [K]']
p1=4
for i in range(p1): 
    
    plt.subplot(p1,1,i+1)
    plt.plot(y_rnn_model_unscaled[:,i],'b')
    plt.plot(y_test_MSAP_unscaled[:,i],'m')
    plt.ylabel("Tempstep " + str(i+1))
    plt.ylabel(str_output[i])
    plt.legend(['Single step','multi step'])
    plt.grid()
    plt.xlabel("Timestep")
    if i == 0:
        plt.title('Multiple step ahead prediction')

# main plot
plt.figure(4)
plt.rcParams["figure.figsize"] = [5.00,3.00]
plt.rcParams["figure.autolayout"] = True
plt.plot(y_test_unscaled[:,0],'r',label='Actual') 
plt.plot(predicted_trajectory_unscaled[:,0],'--g' ,label='Single_step_ahead_pred',linewidth=2)  
plt.plot(y_test_MSAP_unscaled[:,0],':b' ,label='Multiple_step_ahead_pred', linewidth=2) 
plt.legend(loc='best')
plt.xlabel('Time step')
plt.ylabel("$X_{B3}$")
plt.legend(loc='best')
plt.show()

   

# For RMSE calculation a different sets of testing data are used to report the average RMSE values
# The given RMSE is a sample calculation
    
rmse_SAP1=np.zeros(number_of_states)
for i in range(p1):
    rmse_SAP1[i]=rmse(predicted_trajectory_unscaled[:,i],y_test_unscaled[:,i])
print('rmse of model',sum(rmse_SAP1)/len(rmse_SAP1))

rmse_SAP=rmse(predicted_trajectory_unscaled[:600,0],y_test_unscaled[:600,0])
print('rmse_SAP',rmse_SAP)

rmse_MSAP=rmse(y_test_MSAP_unscaled[:600,0],y_test_unscaled[:600,0])
print('rmse_MSAP',rmse_MSAP)


 
## Save the weights and obtain the mathematical expression of the LSTM 

# np.savetxt('./Saved Models3/W_1.txt', W_1)
# np.savetxt('./Saved Models3/W_2.txt', W_2)
# np.savetxt('./Saved Models3/U_1.txt', U_1)
# np.savetxt('./Saved Models3/U_2.txt', U_2)
# np.savetxt('./Saved Models3/b_1.txt', b_1)
# np.savetxt('./Saved Models3/b_2.txt', b_2)
# np.savetxt('./Saved Models3/V_m.txt', V_m)
# np.savetxt('./Saved Models3/b_m.txt', b_m)

# np.savetxt('./Saved Models3/data_min.txt', data_min)
# np.savetxt('./Saved Models3/data_max.txt', data_max)

# np.savetxt('./Saved Models3/predicted_trajectory_unscaled.txt', predicted_trajectory_unscaled)
# np.savetxt('./Saved Models3/y_test_MSAP_unscaled.txt', y_test_MSAP_unscaled)


# np.savetxt('./Saved Models3/y_test_unscaled.txt', y_test_unscaled)
# np.savetxt('./Saved Models3/x_hat_unscaled.txt', x_hat_unscaled)

# np.savetxt("./Saved Models3/y_test.txt", y_test)
# np.savetxt("./Saved Models3/y_test_unscaled.txt", y_test_unscaled)
# np.savetxt("./Saved Models3/X_test.txt", X_test.reshape(X_test.shape[0], int(X_test.shape[1]*X_test.shape[2])))


finish = time.perf_counter()
print('\n Finished in time',finish - start,'seconds')